from gensim.models import Word2Vec
from dataset._acl_imdb import AclImdbDataset, read_acl_imdb
import jieba
import numpy as np


# 数据生成类
class Generator:
    def __init__(self, data):
        self.data = data

    def __iter__(self):
        for sentence in self.data:
            yield sentence


# 获取imdb数据
def get_data_neural():
    acl_imdb_file = read_acl_imdb(download=False)
    dataset = AclImdbDataset(tarFile=acl_imdb_file)

    train_data = dataset.train_data[['content', 'label']].sample(frac=1.0)
    test_data = dataset.test_data[['content', 'label']].sample(frac=1.0)

    x_train = list(train_data['content'])
    y_train = train_data['label']

    x_test = list(test_data['content'])
    y_test = test_data['label']
    train_label = []
    test_label = []
    for item in y_train:
        if item == 'neg':
            train_label.append([1, 0])
        else:
            train_label.append([0, 1])
    for item in y_test:
        if item == 'neg':
            test_label.append([1, 0])
        else:
            test_label.append([0, 1])
    train_label = np.array(train_label)
    test_label = np.array(test_label)
    return x_train, train_label, x_test, test_label


def get_data_machine():
    acl_imdb_file = read_acl_imdb(download=False)
    dataset = AclImdbDataset(tarFile=acl_imdb_file)

    train_data = dataset.train_data[['content', 'label']].sample(frac=1.0)
    test_data = dataset.test_data[['content', 'label']].sample(frac=1.0)

    x_train = list(train_data['content'])
    y_train = train_data['label']

    x_test = list(test_data['content'])
    y_test = test_data['label']
    train_label = []
    test_label = []
    for item in y_train:
        if item == 'neg':
            train_label.append(0)
        else:
            train_label.append(1)
    for item in y_test:
        if item == 'neg':
            test_label.append(0)
        else:
            test_label.append(1)
    train_label = np.array(train_label)
    test_label = np.array(test_label)
    return x_train, train_label, x_test, test_label

# 去停用词并将数据统一数据长度
def remove_stopword(data):
    stop_word_list = ['very', 'ourselves', 'am', 'doesn', 'through', 'me', 'against', 'up', 'just', 'her', 'ours',
                      'couldn', 'because', 'is', 'isn', 'it', 'only', 'in', 'such', 'too', 'mustn', 'under', 'their',
                      'if', 'to', 'my', 'himself', 'after', 'why', 'while', 'can', 'each', 'itself', 'his', 'all',
                      'once',
                      'herself', 'more', 'our', 'they', 'hasn', 'on', 'ma', 'them', 'its', 'where', 'did', 'll', 'you',
                      'didn', 'nor', 'as', 'now', 'before', 'those', 'yours', 'from', 'who', 'was', 'm', 'been', 'will',
                      'into', 'same', 'how', 'some', 'of', 'out', 'with', 's', 'being', 't', 'mightn', 'she', 'again',
                      'be',
                      'by', 'shan', 'have', 'yourselves', 'needn', 'and', 'are', 'o', 'these', 'further', 'most',
                      'yourself', ',', '，', '.', '。', '!', '~', '`', '@', '#', '$', '%', '^', '&', '*', '(', ')', '-',
                      '+',
                      '=', '-', 'having', 'are', 'here', 'he', 'were', 'but', 'this', 'myself', 'own', 'we', 'so', 'i',
                      'does', 'both',
                      'when', 'between', 'd', 'had', 'the', 'y', 'has', 'down', 'off', 'than', 'haven', 'whom',
                      'wouldn',
                      'should', 've', 'over', 'themselves', 'few', 'then', 'hadn', 'what', 'until', 'won', 'no',
                      'about',
                      'any', 'that', 'for', 'shouldn', 'don', 'do', 'there', 'doing', 'an', 'or', 'ain', 'hers', 'wasn',
                      'weren', 'above', 'a', 'at', 'your', 'theirs', 'below', 'other', 'not', 're', 'him', 'during',
                      'which', '.........', ' ', "'", ':', '"', '<', '/', 'br', '>', '?', '...', '\x85', '1', '2', '3',
                      '4',
                      '5', '6', '7', '8', '9', '10', ';', '\x97']
    result_all = []
    for i in range(len(data)):
        result = [k for k in jieba.lcut(data[i], cut_all=False) if k not in stop_word_list]
        if len(result) > 200:
            result_all.append(result[:200])
        else:
            result.extend([None] * (200 - len(result)))
            result_all.append(result)
    return result_all


# 训练word2vec模型并保存
def train_word2vec():
    x_train, train_label, x_test, test_label = get_data()
    x_train = remove_stopword(x_train)
    x_test = remove_stopword(x_test)
    x_train.extend(x_test)
    sentence_genertor = Generator(x_train)
    train_model = Word2Vec(sentence_genertor, vector_size=100, window=5, min_count=1, workers=4)
    train_model.save('./word2vec_l')


# 使用word2vec模型
def use_word2vec(data):
    model = Word2Vec.load('./word2vec_l')
    all_result = []
    for word_list in data:
        temp_list = []
        for word in word_list:
            temp_list.append(model.wv.get_vector(word))
        all_result.append(temp_list)
    return all_result


# #
# if __name__ == '__main__':
#     # train_word2vec()
